import torch
from torch.utils.data import DataLoader
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from sklearn.metrics import confusion_matrix

from model.mymodel import *
from data.dataset import TFs_Dataset  # 你的数据集类

def remove_module_prefix(state_dict):
    """去掉state_dict里key的'module.'前缀，适应DataParallel保存的权重"""
    new_state_dict = {}
    for k, v in state_dict.items():
        if k.startswith('module.'):
            new_key = k[7:]
        else:
            new_key = k
        new_state_dict[new_key] = v
    return new_state_dict

def load_model(checkpoint_path, device, num_classes=25):
    model = ResNet18(num_classes=num_classes)
    checkpoint = torch.load(checkpoint_path, map_location=device)
    state_dict = checkpoint['state_dict']
    state_dict = remove_module_prefix(state_dict)
    model.load_state_dict(state_dict)
    model.to(device)
    model.eval()
    return model

def plot_confusion_matrix_with_acc(y_true, y_pred, class_names, save_path):
    cm = confusion_matrix(y_true, y_pred, labels=range(len(class_names)))
    per_class_acc = cm.diagonal() / (cm.sum(axis=1) + 1e-10)  # 防止除零

    plt.figure(figsize=(14, 10))
    sns.heatmap(cm, annot=True, fmt='d', cmap='Blues',
                xticklabels=class_names, yticklabels=class_names)

    plt.xlabel('Predicted Label')
    plt.ylabel('True Label')
    plt.title('Confusion Matrix with Per-Class Accuracy')

    # 在混淆矩阵右侧显示每类准确率
    for i, acc in enumerate(per_class_acc):
        plt.text(len(class_names) + 0.5, i + 0.5, f'{acc*100:.2f}%', va='center', ha='left', fontsize=10)

    plt.xlim(0, len(class_names) + 2)  # 给右侧文字留位置
    plt.tight_layout()
    plt.savefig(save_path)
    plt.close()
    print(f"[INFO] 混淆矩阵保存到 {save_path}")

def main():
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    print(f"[INFO] 使用设备: {device}")

    # 路径和参数，按实际修改
    data_path = '/media/ubuntu/Data/huyanlong/ZJU_NPY_50ms'  
    checkpoint_path = '/home/ubuntu/huyl/Freshman_Guidance/output/Resnet18_lr0.1/checkpoint_dir/Resnet18-model_best_epoch_34.pth'
    batch_size = 16
    num_classes = 25
    save_path = '/home/ubuntu/huyl/Freshman_Guidance/confusion_matrix_ResNet18_0.1.png'

    # 加载数据集（用你现成的TFs_Dataset）
    dataset = TFs_Dataset(data_path, select_class='all')
    dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=False, num_workers=4)

    # 使用你提供的标签映射生成类别名称
    label_map = {
        'T0000': 0, 'T0001': 1, 'T0010': 2, 'T0011': 3, 'T0100': 4,
        'T0101': 5, 'T0110': 6, 'T0111': 7, 'T1000': 8, 'T1001': 9,
        'T1010': 10, 'T1011': 11, 'T1100': 12, 'T1101': 13, 'T1110': 14,
        'T1111': 15, 'T10000': 16, 'T10001': 17, 'T10010': 18, 'T10011': 19,
        'T10100': 20, 'T10101': 21, 'T10110': 22, 'T10111': 23, 'T11000': 24
    }

    # 生成类别名称列表
    class_names = [key for key, _ in sorted(label_map.items(), key=lambda item: item[1])]

    # 加载模型
    model = load_model(checkpoint_path, device, num_classes=num_classes)

    all_labels = []
    all_preds = []

    with torch.no_grad():
        for imgs, labels in dataloader:
            imgs = imgs.to(device)
            labels = labels.to(device)

            outputs = model(imgs)
            preds = outputs.argmax(dim=1)

            all_labels.extend(labels.cpu().numpy())
            all_preds.extend(preds.cpu().numpy())

    plot_confusion_matrix_with_acc(all_labels, all_preds, class_names, save_path)

if __name__ == '__main__':
    main()
